package io.kiki.sba.registry.server.clientInterface.push;

import io.kiki.sba.registry.util.ConcurrentUtils;
import io.kiki.sba.registry.util.SystemUtils;
import io.kiki.sba.registry.util.WakeUpLoopExecuteTask;
import lombok.Setter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.net.InetSocketAddress;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

import static io.kiki.sba.registry.server.clientInterface.push.PushMetrics.Push.*;

public final class PushTaskWorkers {
    private static final Logger logger = LoggerFactory.getLogger(PushTaskWorkers.class);
    private static final String KEY_MAX_BUFFERED_SIZE = "registry.session.push_task.max_buffered";
    private static final int MAX_BUFFERED_SIZE = SystemUtils.getSystemInteger(KEY_MAX_BUFFERED_SIZE, 10000);

    final Worker[] workers;
    private PushEfficiencyImproveConfig pushEfficiencyImproveConfig;

    PushTaskWorkers(int workerSize) {
        this.workers = new Worker[workerSize];
        for (int i = 0; i < workerSize; i++) {
            Worker worker = new Worker();
            this.workers[i] = worker;
            ConcurrentUtils.createDaemonThread("PushTaskBuffer-" + i, worker).start();
        }
    }

    public void setPushTaskWorkWaitingMillis(int workWaitingMillis) {
        for (Worker worker : workers) {
            worker.setWaitingMillis(workWaitingMillis);
        }
    }

    boolean add(PushTask pushTask) {
        final PushTaskKey pushTaskKey = pushTaskKey(pushTask);
        final Worker worker = workerOf(pushTaskKey);
        boolean result = trtAdd(pushTask);
        if (null != pushEfficiencyImproveConfig && pushEfficiencyImproveConfig.fetchPushTaskWake(pushTask.subscriber.getAppName())) {
            worker.wakeup();
        }
        return result;
    }

    boolean trtAdd(PushTask pushTask) {
        final PushTaskKey pushTaskKey = pushTaskKey(pushTask);
        final Worker worker = workerOf(pushTaskKey);
        if (worker.map.size() >= MAX_BUFFERED_SIZE) {
            return false;
        }
        if (worker.map.putIfAbsent(pushTaskKey, pushTask) == null) {
            // fast path
            wakeup(worker, pushTask);
            BUFFER_NEW_COUNTER.inc();
            return true;
        }

        for (; ; ) {
            final PushTask prev = worker.map.get(pushTaskKey);
            if (prev == null) {
                if (worker.map.putIfAbsent(pushTaskKey, pushTask) == null) {
                    // prev has remove at this time
                    wakeup(worker, pushTask);
                    BUFFER_NEW_COUNTER.inc();
                    return true;
                }
            } else {
                if (pushTask.afterThan(prev)) {
                    final long originExpireTimestamp = pushTask.expireTimestamp;
                    // update the expireTimestamp as prev's, avoid the push block by the continues fire
                    pushTask.expireTimestamp = prev.expireTimestamp;
                    if (worker.map.replace(pushTaskKey, prev, pushTask)) {
                        wakeup(worker, pushTask);
                        BUFFER_REPLACE_COUNTER.inc();
                        return true;
                    } else {
                        // put failed, recover the expireTimestamp
                        pushTask.expireTimestamp = originExpireTimestamp;
                    }
                } else {
                    BUFFER_SKIP_COUNTER.inc();
                    logger.info("[SkipBuffer]key={},prev={},ver={}/{},now={},ver={}/{},retry={}", pushTaskKey, prev.traceId, prev.dataCenterToSubDatumMap.getDataCenterToSubDatumVersionMap(), prev.trace.cause.type, pushTask.traceId, pushTask.dataCenterToSubDatumMap.getDataCenterToSubDatumVersionMap(), prev.trace.cause.type, pushTask.retryCount);
                    return false;
                }
            }
        }
    }

    private void wakeup(Worker worker, PushTask pushTask) {
        if (pushTask.trace.cause.type.noDelay) {
            worker.wakeup();
        }
    }

    int commitPushTask(Worker worker) {
        int bufferedSize = worker.map.size();
        if (bufferedSize >= MAX_BUFFERED_SIZE) {
            logger.warn("arrived max buffered size: buffered={}", bufferedSize);
        }

        List<PushTask> pushTaskList = worker.transferAndMerge();
        int count = 0;
        for (PushTask pushTask : pushTaskList) {
            if (pushTask.commit()) {
                count++;
            }
        }
        if (pushTaskList.size() > 0 || count > 0) {
            logger.info("buffers={},commits={}", pushTaskList.size(), count);
        }
        return count;
    }

    PushTaskKey pushTaskKey(PushTask pushTask) {
        return new PushTaskKey(pushTask.dataCenterToSubDatumMap.dataCenterSet(), pushTask.pushingTaskKey.inetSocketAddress, pushTask.subscriber.getDataInfoId(), pushTask.registerIdToSubscriberMap.keySet());
    }

    private Worker workerOf(Object key) {
        int n = (key.hashCode() & 0x7fffffff) % workers.length;
        return workers[n];
    }

    public int size() {
        int size = 0;
        for (Worker w : workers) {
            size += w.map.size();
        }
        return size;
    }


    void suspend() {
        for (Worker w : workers) {
            w.suspend();
        }
    }


    void resume() {
        for (Worker worker : workers) {
            worker.resume();
        }
    }

    static final class PushTaskKey {
        final Set<String> dataCenterSet;
        final String dataInfoId;
        final InetSocketAddress inetSocketAddress;
        final Set<String> registerIdSet;
        final int hashCode;

        PushTaskKey(Set<String> dataCenterSet, InetSocketAddress inetSocketAddress, String dataInfoId, Set<String> registerIdSet) {
            // all data change push task dataCenters.size=1
            if (dataCenterSet.size() > 1) {
                this.dataCenterSet = new TreeSet<>(dataCenterSet);
            } else {
                this.dataCenterSet = dataCenterSet;
            }
            this.dataInfoId = dataInfoId;
            this.inetSocketAddress = inetSocketAddress;
            this.registerIdSet = registerIdSet;
            if (registerIdSet.size() <= 1) {
                this.hashCode = Objects.hash(dataCenterSet, inetSocketAddress, dataInfoId, registerIdSet);
            } else {
                // sort the subscriberIds
                this.hashCode = Objects.hash(dataCenterSet, inetSocketAddress, dataInfoId, new TreeSet<>(registerIdSet));
            }
        }

        @Override
        public boolean equals(Object o) {
            if (this == o)
                return true;
            if (o == null || getClass() != o.getClass())
                return false;
            PushTaskKey that = (PushTaskKey) o;
            return hashCode == that.hashCode && Objects.equals(dataInfoId, that.dataInfoId) && Objects.equals(inetSocketAddress, that.inetSocketAddress) && Objects.equals(dataCenterSet, that.dataCenterSet) && Objects.equals(registerIdSet, that.registerIdSet);
        }

        @Override
        public int hashCode() {
            return hashCode;
        }

    }

    final class Worker extends WakeUpLoopExecuteTask {
        final Map<PushTaskKey, PushTask> map = new ConcurrentHashMap<>(4096);
        @Setter
        private int waitingMillis = 200;

        @Override
        public void _execute_() {
            commitPushTask(this);
        }

        @Override
        public int getWaitingMillis() {
            return waitingMillis;
        }

        private List<PushTask> transferAndMerge() {
            if (map.isEmpty()) {
                return Collections.emptyList();
            }
            List<PushTask> list = new ArrayList<>(1024);
            final long now = System.currentTimeMillis();
            for (Map.Entry<PushTaskKey, PushTask> entry : map.entrySet()) {
                final PushTask pushTask = entry.getValue();
                // no delay or expire, push immediately
                if (pushTask.trace.cause.type.noDelay || pushTask.expireTimestamp <= now) {
                    list.add(pushTask);
                    // the task maybe update
                    map.remove(entry.getKey(), pushTask);
                }
            }
            return list;
        }
    }
}
